/*
  Copyright (C) 2011, Kyungjoo Kim
  
  This file is part of LINAL (LINear ALgebra)
  
  All rights reserved.

  This library is free software; you can redistribute it and/or
  modify it under the terms of the GNU Lesser General Public
  License as published by the Free Software Foundation; either
  version 2.1 of the License, or (at your option) any later version.
  
  This library is distributed in the hope that it will be useful,
  but WITHOUT ANY WARRANTY; without even the implied warranty of
  MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
  Lesser General Public License for more details.
  
  You should have received a copy of the GNU Lesser General Public
  License along with this library; if not, write to the Free Software
  Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA  02110-1301  USA
  Also add information on how to contact you by electronic and paper mail.
  
  
  Kyungjoo Kim
  iamkyungjoo@gmail.com
*/
#include "linal/common.hxx"
#include "linal/util.hxx"

#include "linal/format.hxx"

#include "linal/operation.hxx"

#ifdef LINAL_DISPATCHER_ENABLE
#include "linal/device/operation.hxx"
#include "linal/device/device.hxx"
#include "linal/device/cache.hxx"
#include "linal/dispatcher.hxx"
#else
#include "linal/device/host/operation.hxx"
#endif

namespace linal {

  static int l_t_upper_trsm   ( int diag, int trans, Hier_ A, Hier_ B );
  static int l_t_lower_trsm   ( int diag, int trans, Hier_ A, Hier_ B );
  static int l_nt_upper_trsm  ( int diag, int trans, Hier_ A, Hier_ B );
  static int l_nt_lower_trsm  ( int diag, int trans, Hier_ A, Hier_ B );
  static int r_t_upper_trsm   ( int diag, int trans, Hier_ A, Hier_ B );
  static int r_t_lower_trsm   ( int diag, int trans, Hier_ A, Hier_ B );
  static int r_nt_upper_trsm  ( int diag, int trans, Hier_ A, Hier_ B );
  static int r_nt_lower_trsm  ( int diag, int trans, Hier_ A, Hier_ B );
  static int trsm_update      ( int side, int uplo,
                                int trans, int diag,
                                Hier_ A, Hier_ B );

  int trsm( int side, int uplo, int trans, int diag,
            FLA_Obj alpha, Hier_ A, Hier_ B ) {

    LINAL_ERROR_DETAIL( check_side     ( side  ), 
                        "%s\n %s %d\n",
                        LINAL_ERROR_CHECK_SIDE,
                        "side = ", side );
    LINAL_ERROR_DETAIL( check_uplo     ( uplo  ), 
                        "%s\n %s %d\n",
                        LINAL_ERROR_CHECK_UPLO,
                        "uplo = ", uplo );
    LINAL_ERROR_DETAIL( check_trans_all( trans ),
                        "%s\n %s %d\n",
                        LINAL_ERROR_CHECK_TRANS,
                        "trans = ", trans );
    LINAL_ERROR_DETAIL( check_diag     ( diag  ), 
                        "%s\n %s %d\n",
                        LINAL_ERROR_CHECK_DIAG,
                        "diag = ", diag );

    if ( !A.get_m() || !A.get_n() || 
         !B.get_m() || !B.get_n() ) return true;

    // Scalar alpha should be applied first otherwise sign is floping.
    scal( alpha, B );

    switch (uplo) {
    case FLA_UPPER_TRIANGULAR:
      {
        switch (side) {
        case FLA_LEFT:
          {
            switch (trans) {
            case FLA_NO_TRANSPOSE:
            case FLA_CONJ_NO_TRANSPOSE:
              {
                l_nt_upper_trsm( diag, trans, A, B ); 
                break;
              }
            case FLA_TRANSPOSE:
            case FLA_CONJ_TRANSPOSE:
              {
                l_t_upper_trsm ( diag, trans, A, B ); 
                break;
              }
            }
            break;
          }
        case FLA_RIGHT:
          {
            switch (trans) {
            case FLA_NO_TRANSPOSE:
            case FLA_CONJ_NO_TRANSPOSE:
              {
                r_nt_upper_trsm( diag, trans, A, B ); 
                break;
              }
            case FLA_TRANSPOSE:
            case FLA_CONJ_TRANSPOSE:
              {
                r_t_upper_trsm ( diag, trans, A, B);  
                break;
              }
            } 
            break;
          }
        }
        break;
      }
    case FLA_LOWER_TRIANGULAR:
      {
        switch (side) {
        case FLA_LEFT:
          {
            switch (trans) {
            case FLA_NO_TRANSPOSE:
            case FLA_CONJ_NO_TRANSPOSE:
              {
                l_nt_lower_trsm(diag, trans, A, B);
                break;
              }
            case FLA_TRANSPOSE:
            case FLA_CONJ_TRANSPOSE:
              {
                l_t_lower_trsm ( diag, trans, A, B ); 
                break;
              }
            }
            break;
          }
        case FLA_RIGHT:
          {
            switch (trans) {
            case FLA_NO_TRANSPOSE:
            case FLA_CONJ_NO_TRANSPOSE:
              {
                r_nt_lower_trsm( diag, trans, A, B ); 
                break;
              }
            case FLA_TRANSPOSE:
            case FLA_CONJ_TRANSPOSE:
              {
                r_t_lower_trsm ( diag, trans, A, B);  
                break;
              }
            } 
            break;
          }
        }
        break;
      }
    }

    return true;
  }

  int l_nt_upper_trsm( int diag, int trans,
                       Hier_ A, Hier_ B) {
    FLA_Obj 
      ATL, ATR, ABL, ABR,
      A00, A01,  A02, A10, A11, A12, A20, A21, A22,
      BT, BB, BL, BR,
      B0, B1, B2;

    LINAL_TR_LEFT(FLA_BR, FLA_BOTTOM);

    while ( FLA_Obj_length( ABR ) < FLA_Obj_length( ~A ) ){
      LINAL_TR_LEFT_LOOP_HEAD(FLA_TL, FLA_TOP);
      //------------------------------------------------------------
      trsm_update( FLA_LEFT,
                   FLA_UPPER_TRIANGULAR,
                   FLA_NO_TRANSPOSE,
                   diag,
                   Hier_(A11), Hier_(B1) );

      gemm( FLA_NO_TRANSPOSE, FLA_NO_TRANSPOSE,
            FLA_MINUS_ONE, Hier_(A01), Hier_(B1),
            FLA_ONE, Hier_(B0) );
      //------------------------------------------------------------
      LINAL_TR_LEFT_LOOP_TAIL(FLA_BR, FLA_BOTTOM);
    }
    return true;
  }
  
  int l_nt_lower_trsm( int diag, int trans,
                       Hier_ A, Hier_ B) {
    FLA_Obj 
      ATL, ATR, ABL, ABR,
      A00, A01,  A02, A10, A11, A12, A20, A21, A22,
      BT, BB, BL, BR,
      B0, B1, B2;

    LINAL_TR_LEFT(FLA_TL, FLA_TOP);

    while ( FLA_Obj_length( ATL ) < FLA_Obj_length( ~A ) ){
      LINAL_TR_LEFT_LOOP_HEAD(FLA_BR, FLA_BOTTOM);
      //------------------------------------------------------------
      trsm_update( FLA_LEFT,
                   FLA_LOWER_TRIANGULAR,
                   FLA_NO_TRANSPOSE,
                   diag,
                   Hier_(A11), Hier_(B1) );

      gemm( FLA_NO_TRANSPOSE, FLA_NO_TRANSPOSE,
            FLA_MINUS_ONE, Hier_(A21), Hier_(B1),
            FLA_ONE, Hier_(B2) );
      //------------------------------------------------------------
      LINAL_TR_LEFT_LOOP_TAIL(FLA_TL, FLA_TOP);
    }

    return true;
  }
  int l_t_upper_trsm( int diag, int trans,
                      Hier_ A, Hier_ B ) {
    FLA_Obj 
      ATL, ATR, ABL, ABR,
      A00, A01,  A02, A10, A11, A12, A20, A21, A22,
      BT, BB, BL, BR,
      B0, B1, B2;

    LINAL_TR_LEFT(FLA_TL, FLA_TOP);

    while ( FLA_Obj_length( ATL ) < FLA_Obj_length( ~A ) ){
      LINAL_TR_LEFT_LOOP_HEAD(FLA_BR, FLA_BOTTOM);
      //------------------------------------------------------------
      trsm_update( FLA_LEFT,
                   FLA_UPPER_TRIANGULAR,
                   trans,
                   diag,
                   Hier_(A11), Hier_(B1) );
      
      gemm( trans, FLA_NO_TRANSPOSE,
            FLA_MINUS_ONE, Hier_(A12), Hier_(B1),
            FLA_ONE, Hier_(B2) );
      //------------------------------------------------------------
      LINAL_TR_LEFT_LOOP_TAIL(FLA_TL, FLA_TOP);
    }

    return true;
  }
  int l_t_lower_trsm( int diag, int trans,
                      Hier_ A, Hier_ B ) {
    FLA_Obj 
      ATL, ATR, ABL, ABR,
      A00, A01,  A02, A10, A11, A12, A20, A21, A22,
      BT, BB, BL, BR,
      B0, B1, B2;

    LINAL_TR_LEFT(FLA_BR, FLA_BOTTOM);

    while ( FLA_Obj_length( ABR ) < FLA_Obj_length( ~A ) ){
      LINAL_TR_LEFT_LOOP_HEAD(FLA_TL, FLA_TOP);
      //------------------------------------------------------------
      trsm_update( FLA_LEFT,
                   FLA_LOWER_TRIANGULAR,
                   trans,
                   diag,
                   Hier_(A11), Hier_(B1) );

      gemm( trans, FLA_NO_TRANSPOSE,
            FLA_MINUS_ONE, Hier_(A10), Hier_(B1),
            FLA_ONE, Hier_(B0) );
      //------------------------------------------------------------
      LINAL_TR_LEFT_LOOP_TAIL(FLA_BR, FLA_BOTTOM);
    }

    return true;
  }

  int r_t_upper_trsm( int diag, int trans,
                      Hier_ A, Hier_ B ) {
    FLA_Obj 
      ATL, ATR, ABL, ABR,
      A00, A01,  A02, A10, A11, A12, A20, A21, A22,
      BT, BB, BL, BR,
      B0, B1, B2;

    LINAL_TR_RIGHT(FLA_BR, FLA_RIGHT);

    while ( FLA_Obj_length( ABR ) < FLA_Obj_length( ~A ) ){
      LINAL_TR_RIGHT_LOOP_HEAD(FLA_TL, FLA_LEFT);
      //------------------------------------------------------------
      trsm_update( FLA_RIGHT,
                   FLA_UPPER_TRIANGULAR,
                   trans,
                   diag,
                   Hier_(A11), Hier_(B1) );

      gemm( FLA_NO_TRANSPOSE, trans,
            FLA_MINUS_ONE, Hier_(B1), Hier_(A01),
            FLA_ONE, Hier_(B0) );
      //------------------------------------------------------------
      LINAL_TR_RIGHT_LOOP_TAIL(FLA_BR, FLA_RIGHT);
    }

    return true;
  }
  int r_t_lower_trsm( int diag, int trans,
                      Hier_ A, Hier_ B) {
    FLA_Obj 
      ATL, ATR, ABL, ABR,
      A00, A01,  A02, A10, A11, A12, A20, A21, A22,
      BT, BB, BL, BR,
      B0, B1, B2;

    LINAL_TR_RIGHT(FLA_TL, FLA_LEFT);

    while ( FLA_Obj_length( ATL ) < FLA_Obj_length( ~A ) ){
      LINAL_TR_RIGHT_LOOP_HEAD(FLA_BR, FLA_RIGHT);
      //------------------------------------------------------------
      trsm_update( FLA_RIGHT,
                   FLA_LOWER_TRIANGULAR,
                   trans,
                   diag,
                   Hier_(A11), Hier_(B1) );

      gemm( FLA_NO_TRANSPOSE, trans,
            FLA_MINUS_ONE, Hier_(B1), Hier_(A21),
            FLA_ONE, Hier_(B2) );
      //------------------------------------------------------------
      LINAL_TR_RIGHT_LOOP_TAIL(FLA_TL, FLA_LEFT);
    }

    return true;
  }
  int r_nt_upper_trsm( int diag, int trans, 
                       Hier_ A, Hier_ B ) {
    FLA_Obj 
      ATL, ATR, ABL, ABR,
      A00, A01,  A02, A10, A11, A12, A20, A21, A22,
      BT, BB, BL, BR,
      B0, B1, B2;

    LINAL_TR_RIGHT(FLA_TL, FLA_LEFT);

    while ( FLA_Obj_length( ATL ) < FLA_Obj_length( ~A ) ){
      LINAL_TR_RIGHT_LOOP_HEAD(FLA_BR, FLA_RIGHT);
      //------------------------------------------------------------
      trsm_update( FLA_RIGHT,
                   FLA_UPPER_TRIANGULAR,
                   FLA_NO_TRANSPOSE,
                   diag,
                   Hier_(A11), Hier_(B1) );

      gemm( FLA_NO_TRANSPOSE, FLA_NO_TRANSPOSE,
            FLA_MINUS_ONE, Hier_(B1), Hier_(A12),
            FLA_ONE, Hier_(B2) );
      //------------------------------------------------------------
      LINAL_TR_RIGHT_LOOP_TAIL(FLA_TL, FLA_LEFT);
    }

    return true;
  }

  int r_nt_lower_trsm( int diag, int trans,
                       Hier_ A, Hier_ B ) {
    FLA_Obj 
      ATL, ATR, ABL, ABR,
      A00, A01,  A02, A10, A11, A12, A20, A21, A22,
      BT, BB, BL, BR,
      B0, B1, B2;

    LINAL_TR_RIGHT(FLA_BR, FLA_RIGHT);

    while ( FLA_Obj_length( ABR ) < FLA_Obj_length( ~A ) ){
      LINAL_TR_RIGHT_LOOP_HEAD(FLA_TL, FLA_LEFT);
      //------------------------------------------------------------
      trsm_update( FLA_RIGHT,
                   FLA_LOWER_TRIANGULAR,
                   FLA_NO_TRANSPOSE,
                   diag,
                   Hier_(A11), Hier_(B1) );
      
      gemm( FLA_NO_TRANSPOSE, FLA_NO_TRANSPOSE,
            FLA_MINUS_ONE, Hier_(B1), Hier_(A10),
            FLA_ONE, Hier_(B0) );
      //------------------------------------------------------------
      LINAL_TR_RIGHT_LOOP_TAIL(FLA_BR, FLA_RIGHT);
    }

    return true;
  }

  int trsm_update( int side, int uplo, 
                   int trans, int diag,
                   Hier_ A, Hier_ B ) {

    // Given input A is one 1x1 block matrix
    // B is column or row block matrix 
    // it apply TRSM on rowise or columnwise for TRSM block
      
    void
      *buffer_A = A.get_buffer(),
      *buffer_B = B.get_buffer();
      
    int
      cs_A = A.get_cs(),
      cs_B = B.get_cs();
      
    for (int k2=0;k2<B.get_n();++k2) {
      for (int k1=0;k1<B.get_m();++k1) {

#pragma omp task firstprivate( k1, k2 ) 
        {
#ifdef LINAL_DISPATCHER_ENABLE
          get_dispatcher()->trsm( side, uplo, trans, diag, 
                                  FLA_ONE, 
                                  elem<FLA_Obj>(buffer_A, cs_A, 1, 0,0), 
                                  elem<FLA_Obj>(buffer_B, cs_B, 1, k1,k2) );

#else
          host::internal::trsm( side, uplo, trans, diag, 
                                FLA_ONE, 
                                elem<FLA_Obj>(buffer_A, cs_A, 1, 0,0), 
                                elem<FLA_Obj>(buffer_B, cs_B, 1, k1,k2) );
#endif
        }
      }
    }

#pragma omp taskwait

    return true;
  }
}
